program automatic test(sm4_io.TB tb_if);

int test_num;
int test_checked;
logic [31:0] payload_key[$];
logic [31:0] payload_plaintext[$];
logic [127:0] pkt2cmp_payload;
logic [127:0] ciphertext_ref;
logic [127:0] plaintext_ref;
logic [127:0] key;

initial begin
    $vcdpluson;
    test_num = 20;
    test_checked = 0;
    reset();
    repeat(test_num) begin
        gen();
        sm4_enc_ref();
        fork
            begin
                key_sch();
                data_enc();
            end
            rev();
        join
        check();
        fork
            data_dec(tb_if.cb.dout);
            rev();
        join
        check();
        test_checked = test_checked+1;        
    end
    repeat(10) @(tb_if.cb);
   end

task  reset();
    tb_if.reset_n <= 1'b0;
    tb_if.cb.key_valid <= 1'b0;
    tb_if.cb.data_valid <= 1'b0;
    tb_if.cb.MK <= 0;
    tb_if.cb.din <= 0;
    tb_if.cb.mode_in <= 2'b10;
    ##2 tb_if.cb.reset_n <= 1'b1;
    repeat(15) @(tb_if.cb);
endtask //reset

task gen();
    payload_key.delete();
    payload_plaintext.delete();
    repeat(4) begin
        payload_key.push_back($urandom);
        payload_plaintext.push_back($urandom);
    end
    plaintext_ref = {payload_plaintext[0],payload_plaintext[1],payload_plaintext[2],payload_plaintext[3]};
    key = {payload_key[0],payload_key[1],payload_key[2],payload_key[3]};
endtask //gen

task send();
    key_sch();
    data_enc();
    data_dec(tb_if.cb.dout);
endtask


task key_sch();
    tb_if.cb.key_valid <= 1'b1;
    tb_if.cb.mode_in <= 2'b00;
    tb_if.cb.MK <= {payload_key[0],payload_key[1],payload_key[2],payload_key[3]};
    while(tb_if.cb.key_ex_finished != 1'b1)
        @(tb_if.cb);
    tb_if.cb.key_valid <= 1'b0;
endtask

task data_enc();
    tb_if.cb.data_valid <= 1'b1;
    tb_if.cb.mode_in <= 2'b11;
    tb_if.cb.din <= {payload_plaintext[0],payload_plaintext[1],payload_plaintext[2],payload_plaintext[3]};
    @(tb_if.cb);
    while(tb_if.cb.data_ready != 1'b1)
        @(tb_if.cb);
    tb_if.cb.data_valid <= 1'b0;
endtask

task data_dec (logic [127:0] ciphertext);
    tb_if.cb.data_valid <= 1'b1;
    tb_if.cb.mode_in <= 2'b01;
    tb_if.cb.din <= ciphertext;
    @(tb_if.cb);
    while(tb_if.cb.data_ready != 1'b1)
        @(tb_if.cb);
    tb_if.cb.data_valid <= 1'b0;
endtask

task rev ();
    get_payload();
endtask //rev

task get_payload ();
    // fork
    //     begin              //add the fork time out process to prevent simutlation run-away when data_ready is not valid
    //         fork
    //             @(posedge tb_if.cb.data_ready);
    //             begin
    //                 repeat(1000) @(tb_if.cb);
    //                 $display("\n%m\n[ERROR]%t data_ready signal timed out !!\n",$realtime);
    //                 $finished;
    //             end
    //         join_any
    //         disable fork;
    //     end
    // join
    @(posedge tb_if.cb.data_ready);
    pkt2cmp_payload = tb_if.cb.dout;
endtask //get_payload

task sm4_enc_ref(); //sm4 encrypto reference model
    static bit [7:0] sbox[256] = { 
    8'hd6,8'h90,8'he9,8'hfe,8'hcc,8'he1,8'h3d,8'hb7,8'h16,8'hb6,8'h14,8'hc2,8'h28,8'hfb,8'h2c,8'h05,  
    8'h2b,8'h67,8'h9a,8'h76,8'h2a,8'hbe,8'h04,8'hc3,8'haa,8'h44,8'h13,8'h26,8'h49,8'h86,8'h06,8'h99,  
    8'h9c,8'h42,8'h50,8'hf4,8'h91,8'hef,8'h98,8'h7a,8'h33,8'h54,8'h0b,8'h43,8'hed,8'hcf,8'hac,8'h62,  
    8'he4,8'hb3,8'h1c,8'ha9,8'hc9,8'h08,8'he8,8'h95,8'h80,8'hdf,8'h94,8'hfa,8'h75,8'h8f,8'h3f,8'ha6,  
    8'h47,8'h07,8'ha7,8'hfc,8'hf3,8'h73,8'h17,8'hba,8'h83,8'h59,8'h3c,8'h19,8'he6,8'h85,8'h4f,8'ha8,  
    8'h68,8'h6b,8'h81,8'hb2,8'h71,8'h64,8'hda,8'h8b,8'hf8,8'heb,8'h0f,8'h4b,8'h70,8'h56,8'h9d,8'h35,  
    8'h1e,8'h24,8'h0e,8'h5e,8'h63,8'h58,8'hd1,8'ha2,8'h25,8'h22,8'h7c,8'h3b,8'h01,8'h21,8'h78,8'h87,  
    8'hd4,8'h00,8'h46,8'h57,8'h9f,8'hd3,8'h27,8'h52,8'h4c,8'h36,8'h02,8'he7,8'ha0,8'hc4,8'hc8,8'h9e,  
    8'hea,8'hbf,8'h8a,8'hd2,8'h40,8'hc7,8'h38,8'hb5,8'ha3,8'hf7,8'hf2,8'hce,8'hf9,8'h61,8'h15,8'ha1,  
    8'he0,8'hae,8'h5d,8'ha4,8'h9b,8'h34,8'h1a,8'h55,8'had,8'h93,8'h32,8'h30,8'hf5,8'h8c,8'hb1,8'he3,  
    8'h1d,8'hf6,8'he2,8'h2e,8'h82,8'h66,8'hca,8'h60,8'hc0,8'h29,8'h23,8'hab,8'h0d,8'h53,8'h4e,8'h6f,  
    8'hd5,8'hdb,8'h37,8'h45,8'hde,8'hfd,8'h8e,8'h2f,8'h03,8'hff,8'h6a,8'h72,8'h6d,8'h6c,8'h5b,8'h51,  
    8'h8d,8'h1b,8'haf,8'h92,8'hbb,8'hdd,8'hbc,8'h7f,8'h11,8'hd9,8'h5c,8'h41,8'h1f,8'h10,8'h5a,8'hd8,  
    8'h0a,8'hc1,8'h31,8'h88,8'ha5,8'hcd,8'h7b,8'hbd,8'h2d,8'h74,8'hd0,8'h12,8'hb8,8'he5,8'hb4,8'hb0,  
    8'h89,8'h69,8'h97,8'h4a,8'h0c,8'h96,8'h77,8'h7e,8'h65,8'hb9,8'hf1,8'h09,8'hc5,8'h6e,8'hc6,8'h84,  
    8'h18,8'hf0,8'h7d,8'hec,8'h3a,8'hdc,8'h4d,8'h20,8'h79,8'hee,8'h5f,8'h3e,8'hd7,8'hcb,8'h39,8'h48
    };
    static bit [31:0] fk[4] = {32'hA3B1BAC6, 32'h56AA3350, 32'h677D9197, 32'hB27022DC};
    static bit [31:0] ck[32] = { 
    32'h00070e15, 32'h1c232a31, 32'h383f464d, 32'h545b6269,  
    32'h70777e85, 32'h8c939aa1, 32'ha8afb6bd, 32'hc4cbd2d9,  
    32'he0e7eef5, 32'hfc030a11, 32'h181f262d, 32'h343b4249,  
    32'h50575e65, 32'h6c737a81, 32'h888f969d, 32'ha4abb2b9,  
    32'hc0c7ced5, 32'hdce3eaf1, 32'hf8ff060d, 32'h141b2229,  
    32'h30373e45, 32'h4c535a61, 32'h686f767d, 32'h848b9299,  
    32'ha0a7aeb5, 32'hbcc3cad1, 32'hd8dfe6ed, 32'hf4fb0209,  
    32'h10171e25, 32'h2c333a41, 32'h484f565d, 32'h646b7279
    };
    automatic bit [31:0] key[4]= {payload_key[0],payload_key[1],payload_key[2],payload_key[3]};
    automatic bit [31:0] plaintext[4] = {payload_plaintext[0],payload_plaintext[1],payload_plaintext[2],payload_plaintext[3]} ;
    bit [31:0] rk_gen[32];
    bit [31:0] ciphertext_gen[32];
    bit [31:0] k0,k1,k2,k3,tmp;
    bit [31:0] x0,x1,x2,x3;
    
    k0 = payload_key[0] ^ fk[0];
    k1 = payload_key[1] ^ fk[1];
    k2 = payload_key[2] ^ fk[2];
    k3 = payload_key[3] ^ fk[3];
    x0 =  payload_plaintext[0];
    x1 =  payload_plaintext[1];
    x2 =  payload_plaintext[2];
    x3 =  payload_plaintext[3];

    for(int i=0;i<32;i=i+1) begin
        tmp = k1 ^k2 ^k3 ^ck[i];
        tmp = {sbox[tmp[31:24]],sbox[tmp[23:16]],sbox[tmp[15:8]],sbox[tmp[7:0]]};
        k0 = k0 ^ (tmp ^ {tmp[18:0],tmp[31:19]} ^ {tmp[8:0],tmp[31:9]});
        rk_gen[i] = k0;
        k0 = k1;
        k1 = k2;
        k2 = k3;
        k3 = rk_gen[i];
    end

    for(int i=0;i<32;i=i+1) begin
        tmp = x1 ^ x2 ^ x3 ^ rk_gen[i];
        tmp = {sbox[tmp[31:24]],sbox[tmp[23:16]],sbox[tmp[15:8]],sbox[tmp[7:0]]};
        x0 = x0 ^ (tmp ^ {tmp[29:0],tmp[31:30]} ^ {tmp[21:0],tmp[31:22]} ^ {tmp[13:0],tmp[31:14]} ^ {tmp[7:0],tmp[31:8]});
        ciphertext_gen[i] = x0;
        x0 = x1;
        x1 = x2;
        x2 = x3;
        x3 = ciphertext_gen[i];
    end
    ciphertext_ref = {ciphertext_gen[31],ciphertext_gen[30],ciphertext_gen[29],ciphertext_gen[28]};

endtask

function bit compare(ref string message);
    if((tb_if.mode_in == 3) && (pkt2cmp_payload == ciphertext_ref))
        message = "SM4 Encrypto Successfully Compared!";
    else if(tb_if.mode_in == 1 && pkt2cmp_payload == plaintext_ref)
        message = "SM4 Decrypto Successfully Compared!";
    else begin
        if(tb_if.mode_in == 3) begin
            message = "Encrypto ERROR\n";
            message = {message, $sformatf("key: %32h\n output: %32h\n reference: %32h",key,pkt2cmp_payload,ciphertext_ref)};
            return(0);
        end
        else if(tb_if.mode_in == 1) begin
            message = "Decrypto ERROR\n";
            message = {message, $sformatf("key: %32h\n output: %32h\n reference: %32h",key,pkt2cmp_payload,plaintext_ref)};
            return(0);
        end
    end
    return(1);
endfunction

task check();
    string message;
    if(!compare(message)) begin
        $display("\n%m\n[ERROR]%t test_num #%2d %s\n",$realtime,test_checked,message);
        $finish;
    end
    $display("[NOTE]%t test_num #%2d %s\n",$realtime,test_checked,message);
endtask //check


endprogram
